﻿#region Copyright information
// 
// Copyright © 2005-2013 Yongkee Cho. All rights reserved.
// 
// This code is a part of the Biological Object Library and governed under the terms of the
// GNU Lesser General  Public License (LGPL) version 2.1 which accompanies this distribution.
// For more information on the LGPL, please visit http://bol.codeplex.com/license.
// 
// - Filename: Id3.cs
// - Author: Yongkee Cho
// - Email: yongkeecho@gmail.com
// - Date Created: 2012-12-08 8:26 AM
// - Last Modified: 2013-01-25 3:59 PM
// 
#endregion
using System;
using System.Collections.Generic;
using System.Linq;
using BOL.Linq.Descriptive;
using BOL.Linq.Probability;

namespace BOL.Algorithms.Classification
{
    /// <summary>
    /// Implements an ID3 (Iterative Dichtomiser 3) algorithm (author: J. Ross Quinlan)
    /// </summary>
    public class Id3<TFeature, TClass> : DecisionTreeClassifierBase<TFeature, TClass>
    {
        #region Constructor

        public Id3(IDecisionTreeNode<TClass> rootNode)
            : base(rootNode)
        {
        }

        #endregion

        #region Public method

        public static Id3<TFeature, TClass> Train<TSource>(IEnumerable<TSource> source, Func<TSource, TClass> classSelector, Func<TSource, TFeature[]> featuresSelector)
        {
            if (source == null)
                throw new ArgumentNullException("source");
            if (classSelector == null)
                throw new ArgumentNullException("classSelector");
            if (featuresSelector == null)
                throw new ArgumentNullException("featuresSelector");
            
            var list = source.ToList();
            var featureLength = list.Select(featuresSelector).First().Length;
            var undecomposedFeatureIndices = Enumerable.Range(0, featureLength).ToList();
            var rootNode = new ConditionTreeNode<TClass>();
            Decompose(list, classSelector, featuresSelector, ref undecomposedFeatureIndices, ref rootNode);

            return new Id3<TFeature, TClass>(rootNode);
        }

        private static void Decompose<TSource>(IList<TSource> list, Func<TSource, TClass> classSelector, Func<TSource, TFeature[]> featuresSelector, ref List<int> undecomposedFeatureIndices, ref ConditionTreeNode<TClass> node)
        {
            var total = list.Count();

            // for each undecomposed feature,
            var group = undecomposedFeatureIndices.ToDictionary(i => i, i => list.GroupBy(x => featuresSelector(x)[i]));
            var undecomposeFeatureStatistics = group.ToDictionary(x => x.Key, x => Estimate(x.Value.ToList(), classSelector));
            var entropy = list.GroupBy(classSelector).ToDictionary(x => x.Key, x => x.Count() / (double)total).ShannonEntropy();
            var gain = undecomposeFeatureStatistics.ToDictionary(x => x.Key, x => entropy - x.Value.Aggregate(0.0, (sum, current) => sum + current.Value.Item1 * current.Value.Item2.ShannonEntropy()));
            var argmax = gain.ArgMax();

            // updates current decision tree node
            node.FeatureIndex = argmax;
            node.Expression = x => x;
            node.Condition = x => x;

            // remove feature with highest gain from the list
            undecomposedFeatureIndices.Remove(argmax);

            // adds child nodes
            var maxGainNode = undecomposeFeatureStatistics[argmax];

            if (undecomposedFeatureIndices.Count > 0)
            {
                var keyMin = maxGainNode.ToDictionary(x => x.Key, x => x.Value.Item1).ArgMin();
                node.Add(new ProbabilityTreeNode<TClass>(keyMin, maxGainNode[keyMin].Item2));

                foreach (var key in maxGainNode.Keys.Where(x => !x.Equals(keyMin)))
                {
                    var current = maxGainNode[key];
                    var child = new ConditionTreeNode<TClass>(key);
                    node.Add(child);

                    Decompose(current.Item3, classSelector, featuresSelector, ref undecomposedFeatureIndices, ref child);
                }
            }
            else
                foreach (var key in maxGainNode.Keys)
                    node.Add(new ProbabilityTreeNode<TClass>(key, maxGainNode[key].Item2));
        }

        private static IDictionary<TFeature, Tuple<double, IDictionary<TClass, double>, IList<TSource>>> Estimate<TSource>(IList<IGrouping<TFeature, TSource>> list, Func<TSource, TClass> classSelector)
        {
            var total = list.Sum(x => x.Count());
            return list.ToDictionary(x => x.Key, x => new Tuple<double, IDictionary<TClass, double>, IList<TSource>>(x.Count() / (double)total, ProbabilityTable(x.ToList(), classSelector), x.ToList()));
        }

        private static IDictionary<TClass, double> ProbabilityTable<TSource>(IList<TSource> list, Func<TSource, TClass> classSelector)
        {
            var total = list.Count();
            return list.GroupBy(classSelector).ToDictionary(x => x.Key, x => x.Count() / (double)total);
        }

        #endregion
    }
}
